package de.unibi.comet.fa;

import java.util.HashMap;
import java.util.Map.Entry;

import de.unibi.comet.util.Log;
import de.unibi.comet.util.LogSpace;

/** A Markov additive chain implemented straightforwardly. All information about
 *  states and transition probabilities is stored in arrays. Supports calculations
 *  in logarithmic domain. */
public class MarkovAdditiveChain {
	/** For all j in preimage[i], the transition probability from j to i is >0. */ 
	private int[][] preimage;
	/** Probabilities corresponding to the states listed in preimage. */
	private double[][] preimageProb;
	/** The emission of each state. */
	private int[] emissions;
	/** Current joint probability distribution: dist[state][count]. */
	private double[][] dist;
	private int maxMass;
	/** If true, all probabilities are logarithmic. */
	private boolean logarithmic;
	
	public MarkovAdditiveChain(int[][] preimage, double[][] preimageProb, int[] emissions) {
		if (preimage.length!=preimageProb.length) throw new IllegalArgumentException();
		if (emissions.length!=preimageProb.length) throw new IllegalArgumentException();
		this.preimage = preimage;
		this.preimageProb = preimageProb;
		this.emissions = emissions;
		this.logarithmic=false;
		dist=null;
	}
	
	/** Returns the number of states. */
	public int getStateCount() { return preimage.length; }
	
	/** Set joined distribution of state and count (in that order).
	 *  NOTE: if logarithmic, then a logarithmic distribution has to be given.
	 *  @param dist Has to be a rectangular two dimensional array. 
	 */ 
	public void setDistribution(double[][] dist) {
		if (dist.length!=preimage.length) throw new IllegalArgumentException();
		maxMass=dist[0].length-1;
		this.dist=dist;
	}
	
	/** Update joint distribution. */
	public void step(int steps) {
		Log.getInstance().startTimer();
		// recycle distributions to aviod recreation
		double[][] tmpDist = new double [dist.length][maxMass+1];
		// main loop
		for (int n=0; n<steps; ++n){
			// iteratate over all states
			for (int i=0; i<preimage.length; ++i) {
				// iterate over all masses < maxMass
				for (int m=0; m<maxMass; ++m) {
					// mass of predecessor state
					int mSource = m - emissions[i];
					// iterate over all states in preimage
					double p=logarithmic?Double.NEGATIVE_INFINITY:0.0;
					if (mSource>=0) {
						for (int k=0; k<preimage[i].length; ++k) {
							if (logarithmic) {
								p=LogSpace.logAdd(p, dist[preimage[i][k]][mSource]+preimageProb[i][k]);
							} else {
								p+=dist[preimage[i][k]][mSource]*preimageProb[i][k];
							}
						}
					}
					tmpDist[i][m]=p;
				}
				// handle case m==maxMass
				double p=logarithmic?Double.NEGATIVE_INFINITY:0.0;
				for (int mSource=maxMass-emissions[i]; mSource<=maxMass; ++mSource) {
					// iterate over all states in preimage
					if (mSource>=0) {
						for (int k=0; k<preimage[i].length; ++k) {
							if (logarithmic) {
								p=LogSpace.logAdd(p, dist[preimage[i][k]][mSource]+preimageProb[i][k]);
							} else {
								p+=dist[preimage[i][k]][mSource]*preimageProb[i][k];
							}
						}
					}
				}
				tmpDist[i][maxMass]=p;
			}
			// swap distributions
			double[][] h = dist;
			dist=tmpDist;
			tmpDist=h;
		}
		if (logarithmic) {
			Log.getInstance().stopTimer("Markov additive chain evolution (logarithmic)");
		} else {
			Log.getInstance().stopTimer("Markov additive chain evolution");
		}
	}
	
	/** Calculate the distribution of the waiting time.
	 * @param startStates States to start in. The algorithm will use the restriction
	 *    of the current distribution on these states to start with. 
	 */
	public double[] waitingTime(Iterable<Integer> startStates, Iterable<Integer> finalStates, int steps) {
		if (logarithmic) throw new IllegalStateException("Not implemented in logarithmic domain.");
		double[] waitingTime = new double[steps+1];
		double[] stateDistribution = getStateDistribution();
		// calculate probability of being in startStart
		double pStart = 0.0;
		for (int i : startStates)pStart+=stateDistribution[i];
		if (pStart==0.0) throw new IllegalStateException("Probability of being in a start state is zero.");
		// calculate new state distribution. distStillWaiting[i] is the probability of being in
		// state i and still waiting for the event to come.
		double[] distStillWaiting = new double[stateDistribution.length];
		for (int i : startStates) distStillWaiting[i]=stateDistribution[i]/pStart;
		// array for fast lookup if a state is a final state
		boolean[] stateIsFinal = new boolean[stateDistribution.length];
		for (int i : finalStates) stateIsFinal[i]=true;
		
		// main loop
		for (int step=1; step<steps; ++step) {
			double[] newDistStillWaiting = new double[dist.length];
			for (int targetState=0; targetState<dist.length; ++targetState){
				for (int i=0; i<preimage[targetState].length; ++i) {
					// transition probability
					double transitionProb = preimageProb[targetState][i];
					int sourceState = preimage[targetState][i];
					// if this is a state we wait for, some probability mass leaves the table
					// distStillWaiting and thereby indicates a chance of the waiting to end in 
					// this step
					if (stateIsFinal[targetState]) {
						double p = distStillWaiting[sourceState]*transitionProb;
						waitingTime[step]+=p;
					} else {
						newDistStillWaiting[targetState]+=distStillWaiting[sourceState]*transitionProb;
					}
				}
			}
			distStillWaiting=newDistStillWaiting;
		}
		// calculate the probabitilty that we have to waits steps or more steps
		double pRest = 0.0;
		for (int i=0; i<dist.length; ++i) {
			pRest+=distStillWaiting[i];
		}
		waitingTime[steps]=pRest;
		
		return waitingTime;
	}
	
	/** Helper method needed for the doubling technique. */
	private double[][][] joinDistributions(double[][][] dist1, double[][][] dist2) {
		int stateCount = preimage.length;
		double[][][] resultDist = new double[stateCount][stateCount][maxMass+1];
		for (int state1=0; state1<stateCount; ++state1) {
			for (int state2=0; state2<stateCount; ++state2) {
				// iteration for mass<maxMass
				for (int mass=0; mass<maxMass; ++mass) {
					// calculate value of dists[bit][state1][state2][mass]
					double p = 0.0;
					for (int intermediateMass=0; intermediateMass<=mass; ++intermediateMass) {
						for (int intermediateState=0; intermediateState<stateCount; ++intermediateState) {
							p+=dist1[state1][intermediateState][intermediateMass]*
							   dist2[intermediateState][state2][mass-intermediateMass];
						}
					}
					resultDist[state1][state2][mass]=p;
				}
				// case mass=maxMass
				// calculate value of dists[bit][state1][state2][mass]
				double p = 0.0;
				for (int intermediateMass1=0; intermediateMass1<=maxMass; ++intermediateMass1) {
					for (int intermediateMass2=maxMass-intermediateMass1; intermediateMass2<=maxMass; ++intermediateMass2) {
						for (int intermediateState=0; intermediateState<stateCount; ++intermediateState) {
							p+=dist1[state1][intermediateState][intermediateMass1]*
							   dist2[intermediateState][state2][intermediateMass2];
						}
					}
					resultDist[state1][state2][maxMass]=p;
				}
			}
		}
		return resultDist;
	}
	
	/** Update distribution using the doubling technique. */
	public double[] stepDoubling(int steps) {
		if (logarithmic) throw new IllegalStateException("Not implemented in logarithmic domain.");
		int stateCount = preimage.length;
		int bitsNeeded = (int)(Math.ceil(Math.log(steps) / Math.log(2))+1.0);
		// distributions dist[k][state1][state2][mass_delta] is the probability of
		// being in state2 and having added mass_delta after 2^k steps, given that
		// we started in state1.
		double[][][][] dists = new double[bitsNeeded][][][];
		
		// initialize first distribution
		dists[0] = new double[stateCount][stateCount][maxMass+1];
		for (int targetState=0; targetState<stateCount; ++targetState) {
			int massDelta = emissions[targetState];
			for (int k=0; k<preimage[targetState].length; ++k) {
				int sourceState=preimage[targetState][k];
				dists[0][sourceState][targetState][massDelta]+=preimageProb[targetState][k];
			}
		}
		// create all needed distributions
		for (int bit=1; bit<bitsNeeded; ++bit) {
			if ((1<<bit)>steps) break;
			dists[bit] = joinDistributions(dists[bit-1],dists[bit-1]);
			// do we still need the previous distribution?
			if (((1<<(bit-1))&steps)==0) dists[bit-1]=null;
		}
		// join distributions to create final
		double[][][] targetDist = null;
		for (double[][][] d : dists) {
			if (d==null) continue;
			if (targetDist==null) {
				targetDist=d;
			} else {
				targetDist=joinDistributions(targetDist,d);
			}
		}
		// sum up to calculate mass distribution
		double[] result = new double[maxMass+1];
		for (int startState=0; startState<stateCount; ++startState) {
			for (int startMass=0; startMass<=maxMass; ++startMass) {
				double p0 = dist[startState][startMass]; 
				if (p0==0.0) continue;
				for (int massDiff=0; massDiff<=maxMass; ++massDiff) {
					int m = startMass + massDiff;
					double p1 = 0.0;
					for (int endState=0; endState<stateCount; ++endState) {
						p1+=targetDist[startState][endState][massDiff];
					}
					result[m]+=p0*p1;
				}	
			}
		}
		return result;
	}
	
	/** Converts all probabilities to the logarithmic domain. */
	public void convertToLogSpace() {
		if (logarithmic) throw new IllegalStateException("Already logarithmic!");
		for (int state=0; state<preimageProb.length; ++state) {
			for (int i=0; i<preimageProb[state].length; ++i) {
				preimageProb[state][i]=Math.log(preimageProb[state][i]);
			}
		}
		if (dist!=null) {
			for (int state=0; state<dist.length; ++state) {
				for (int mass=0; mass<dist[state].length; ++mass) {
					dist[state][mass]=Math.log(dist[state][mass]);
				}
			}
		}
		logarithmic=true;
	}
	
	public double[][] getDistribution() {return dist;}
	
	public double getMaxMass() { return maxMass; }

	public double[] getStateDistribution() {
		if (logarithmic) throw new IllegalStateException("Not implemented in logarithmic domain.");
		double[] result = new double[dist.length];
		for (int state=0; state<dist.length; ++state) {
			double p = 0.0;
			for (int mass=0; mass<=maxMass; ++mass) {
				p+=dist[state][mass];
			}
			result[state]=p;
		}
		return result;
	}

	/** Calculates the count distribution by (for each count) summing over all states. */
	public double[] getMassDistribution() {
		double[] result = new double[maxMass+1];
		for (int mass=0; mass<=maxMass; ++mass) {
			double p = logarithmic?Double.NEGATIVE_INFINITY:0.0;
			for (int state=0; state<dist.length; ++state) {
				if (logarithmic) {
					p=LogSpace.logAdd(p, dist[state][mass]);
				} else {
					p+=dist[state][mass];
				}
			}
			result[mass]=p;
		}
		return result;
	}

	/** Returns the probability of having accumulated the given mass. */
	public double getMassProbability(int mass) {
		double result = logarithmic?Double.NEGATIVE_INFINITY:0.0;
		for (int state=0; state<dist.length; ++state) {
			if (logarithmic) {
				result=LogSpace.logAdd(result, dist[state][mass]);
			} else {
				result+=dist[state][mass];
			}
		}
		return result;
	}
	
	/** Joins states according to a given partition, i.e. states that are in
	 *  the same block will be merged. The distribution is updated accordingly.
	 */
	public MarkovAdditiveChain joinStates(Partition partition) {
		if (logarithmic) throw new IllegalStateException("Not implemented in logarithmic domain.");
		if (partition.size()!=emissions.length) {
			throw new IllegalArgumentException("Partition size is not equal to number of states!");
		}
				
		// initialize variable for new MAC
		int newStateCount = partition.blockCount();
		int[][] newPreimage = new int[newStateCount][];
		double[][] newPreimageProb = new double[newStateCount][];
		int[] newEmissions = new int[newStateCount];
		double[][] newDist = new double[newStateCount][maxMass+1];
		
		// old state distribution
		double[] stateDist = getStateDistribution();
		// calculate new state distribution
		double[] newStateDist = new double[newStateCount];
		for (int newState=0; newState<newStateCount; ++newState) {
			double p = 0.0;
			for (int oldState : partition.block(newState)) {
				p+=stateDist[oldState];
			}
			newStateDist[newState]=p;
		}

		// main step: calulate preimages and new joint distribution
		for (int newState=0; newState<newStateCount; ++newState) {
			int firstOldState = -1;
			HashMap<Integer,Double> newPreimageMap = new HashMap<Integer,Double>();
			for (int oldState : partition.block(newState)) {
				if (firstOldState==-1) firstOldState=oldState;
				for (int i=0; i<preimage[oldState].length; ++i) {
					int oldPreimageState = preimage[oldState][i];
					int newPreimageState = partition.getBlockIndex(oldPreimageState);
					if (newPreimageMap.containsKey(newPreimageState)) {
						newPreimageMap.put(newPreimageState, newPreimageMap.get(newPreimageState)
								+ (preimageProb[oldState][i]*stateDist[oldPreimageState])/newStateDist[newPreimageState]);
					} else {
						newPreimageMap.put(newPreimageState, (preimageProb[oldState][i]*stateDist[oldPreimageState])/newStateDist[newPreimageState]);
					}
				}
				// update new distribution
				for (int mass=0; mass<=maxMass; ++mass) {
					newDist[newState][mass]+=dist[oldState][mass];
				}
			}
			
			// create preimage arrays for new state
			newPreimage[newState]=new int[newPreimageMap.size()];
			newPreimageProb[newState]=new double[newPreimageMap.size()];
			int j = 0;
			for (Entry<Integer,Double> entry : newPreimageMap.entrySet()) {
				newPreimage[newState][j]=entry.getKey();
				newPreimageProb[newState][j]=entry.getValue();
				++j;
			}
						
			// new emissions
			newEmissions[newState]=emissions[firstOldState];
		}
		
		MarkovAdditiveChain result = new MarkovAdditiveChain(newPreimage, newPreimageProb, newEmissions);
		result.setDistribution(newDist);
		return result;
	}
}
